import torch
import torch.nn as nn
from torch import Tensor
from torchsummary import summary
# 一个make_layer代表一个结构，比如【3*3 conv，64 3*3 conv，64]结构进行了3次，放在一个make_layer中
# BasicBlock定义每一组操作，例如【3*3 conv，64 3*3 conv，64]这组结构
# 输入特征矩阵的channels是64，经过一个3x3的卷积核卷积之后，要进行Relu激活函数的激活，
# 再经过一个3x3的卷积核进行卷积，但是在这没有经过激活函数进行激活。
# 将这个输出特征和捷径分支带回来的原始输入，进行相加，而不是叠加或者拼接，
class BasicBlock(nn.Module):
    # 在ResNet18Layer,ResNet34Layer网络结构中，每一组的通道相同。
    # 但是在ResNet50Layer甚至更深的网络结构中给，每一组的通道不同，例如【1*1,64 3*3,64 1*1,256】
    # 这种残差结构先降维再升维，如果式50layer可以定义为4
    expansion = 1  # 可以理解为初始通道64的膨胀系数.每一个_make_layer最后输出通道和这一层第一卷积后的通道倍数关系
    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super().__init__()
        # 在进行虚线卷积的时候，步长不是1，是2
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        # 每一组里第二次卷积，步长都是1
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        # 进行shortcut，实现shortcut直接相加，虚线的shortcut需要做的封装到downsample中
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        # 为了shortcut捷径分支，先将原始输入x保存起来
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        # 看结构shortcut是虚线还是实现，通过参数downsample不是None，也就是传入参数值
        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class ResNet(nn.Module):
    def __init__(self, block, layers, num_classes=1000):
        super().__init__()
        self.inplanes = 64 #代表通道
        # self.dilation = 1
        # padding是same自己计算padding层
        # bias=False没有偏置，所有卷积没有偏置,参数变少计算快
        # 参数1：输入通道3，RGB.
        # 参数padding=3，(224-7+2p)/2 + 1 = 112
        # bias=False，因为最终会使用BN，不需要有偏置，反向求导不会对偏置求导也会加快速度.
        # ResNet所有卷积没有偏置，全连接有偏置
        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(self.inplanes)
        # 将relu操作得到的值是否需要存储一份，设置True代表替换原有值，节省内存加快速度
        # 没有保留原始数据，无法回溯
        self.relu = nn.ReLU(inplace=True)

        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        # 参数2，输出通道，layers是参数传入的，定义了四层网络的每个BasicBlock执行的次数
        self.layer1 = self._make_layer(block, planes=64, blocks=layers[0])
        self.layer2 = self._make_layer(block, planes=128, blocks=layers[1], stride=2)
        self.layer3 = self._make_layer(block, planes=256, blocks=layers[2], stride=2)
        self.layer4 = self._make_layer(block, planes=512, blocks=layers[3], stride=2)
        # (1, 1) 指定了输出特征图的高度和宽度都为 1。这意味着无论输入特征图的大小如何，经过这个自适应平均池化层后，输出特征图都会被缩减为一个单一的值（在每个通道上）。
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * block.expansion, num_classes)
    # 一个make_layer代表一个结构，比如在ResNet34中，【3*3 conv，64 3*3 conv，64]结构进行了3次，放在一个make_layer中
    # planes是输出特征的通道数  block参数就是 BasicBlock对象
    # blocks代表每个BasickBlock的次数，【3*3 conv，64 3*3 conv，64]结构进行了blocks=3次，
    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        # 当输入特征通道和输出特征通道*n不相等时，定义一个捷径分支
        if stride != 1 or self.inplanes != planes * block.expansion:
            # downsample叫做捷径分支
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        # 第一次和每次不一样，直接是实现开始，所以第一次单独传进去
        layers.append(
            block(self.inplanes, planes, stride, downsample)
        )
        self.inplanes = planes * block.expansion
        # blocks代表每一个basickblock有可能执行的次数，可能是2，3，4，6
        for _ in range(1, blocks):
            layers.append(
                block(self.inplanes, planes)
            )

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1) # 展平
        x = self.fc(x) # 全连接层

        return x


def resnet18(num_classes=1000):
    # https://download.pytorch.org/models/resnet18-f37072fd.pth
    return ResNet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes)


def resnet34(num_classes=1000):
    # https://download.pytorch.org/models/resnet34-333f7ec4.pth
    return ResNet(BasicBlock, [3, 4, 6, 3], num_classes=num_classes)


if __name__ == '__main__':
    # model = resnet18(num_classes=3).to("cuda")
    # model = resnet34(num_classes=3).to("cuda")
    model = resnet18(num_classes=3)
    # model = resnet34(num_classes=3)
    # model = resnet50(num_classes=3).to("cuda")
    # fc_inputs = model.fc.in_features
    # model.fc = nn.Sequential(
    #     nn.Linear(fc_inputs, 3),
    #     # nn.ReLU(),
    #     # nn.Dropout(0.4),
    #     # nn.Linear(256, 3),
    # ).to("cuda")
    print(summary(model, (3, 224, 224)))